import pandas as pd
import numpy as np
from typing import List, Dict, Set, Tuple
import os
import matplotlib.pyplot as plt

plt.rcParams["font.family"] = ["sans-serif"]
plt.rcParams["font.sans-serif"] = plt.rcParamsDefault["font.sans-serif"]

def dcg_at_k(relevance: List[float], k: int) -> float:
    """计算DCG@k (Discounted Cumulative Gain)"""
    relevance = np.asarray(relevance)[:k]
    if relevance.size == 0:
        return 0.0
    return relevance[0] + np.sum(relevance[1:] / np.log2(np.arange(2, relevance.size + 1)))

def ndcg_at_k(relevance: List[float], k: int) -> float:
    """计算NDCG@k (Normalized DCG@k)"""
    dcg = dcg_at_k(relevance, k)
    if dcg == 0:
        return 0.0
    idcg = dcg_at_k(sorted(relevance, reverse=True), k)
    return dcg / idcg

def map_at_k(relevance: List[float], k: int) -> float:
    """计算MAP@k (Mean Average Precision@k)"""
    relevance = np.asarray(relevance)[:k]
    if np.sum(relevance) == 0:
        return 0.0
    
    precision_sum = 0.0
    relevant_count = 0
    
    for i, rel in enumerate(relevance):
        if rel == 1:
            relevant_count += 1
            precision_sum += relevant_count / (i + 1)
    
    return precision_sum / np.sum(relevance)

def evaluate_recommendations(
    recommendations: Dict[str, List[str]],
    ground_truth: Dict[str, Set[str]],
    k: int = 10
) -> Tuple[float, float, Dict[str, float], Dict[str, float]]:
    """
    评估推荐结果的NDCG@k和MAP@k
    """
    ndcg_scores = {}
    map_scores = {}
    
    for user_id, rec_list in recommendations.items():
        gt_set = ground_truth.get(user_id, set())
        relevance = [1.0 if item in gt_set else 0.0 for item in rec_list[:k]]
        ndcg_scores[user_id] = ndcg_at_k(relevance, k)
        map_scores[user_id] = map_at_k(relevance, k)
    
    avg_ndcg = np.mean(list(ndcg_scores.values()))
    avg_map = np.mean(list(map_scores.values()))
    
    return avg_ndcg, avg_map, ndcg_scores, map_scores

def load_data(
    recommendations_path: str,
    ratings_path: str,
    user_path: str = None,
    movies_path: str = None
) -> Tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame, pd.DataFrame]:
    """加载所有需要的数据"""
    print("Loading data...")
    
    recommendations = pd.read_csv(recommendations_path)
    print(f"Recommendations data loaded, {len(recommendations)} records")
    
    ratings = pd.read_csv(ratings_path)
    print(f"Ratings data loaded, {len(ratings)} records")
    
    users = None
    if user_path and os.path.exists(user_path):
        users = pd.read_csv(user_path)
        print(f"User data loaded, {len(users)} records")
    
    movies = None
    if movies_path and os.path.exists(movies_path):
        movies = pd.read_csv(movies_path)
        print(f"Movies data loaded, {len(movies)} records")
    
    return recommendations, ratings, users, movies

def prepare_ground_truth(
    ratings: pd.DataFrame,
    recommendations: pd.DataFrame,
    min_rating: float = 4.0
) -> Dict[str, Set[str]]:
    """构建真实相关物品集合"""
    print("Building ground truth set...")
    
    high_rating_mask = ratings['RATING'] >= min_rating
    high_rating_data = ratings[high_rating_mask]
    
    ground_truth = {}
    for _, row in high_rating_data.iterrows():
        user_id = row['USER_MD5']
        movie_id = str(row['MOVIE_ID'])
        if user_id not in ground_truth:
            ground_truth[user_id] = set()
        ground_truth[user_id].add(movie_id)
    
    clicked_mask = recommendations['clicked'] == 1
    clicked_data = recommendations[clicked_mask]
    
    for _, row in clicked_data.iterrows():
        user_id = row['user_id']
        movie_id = str(row['movie_id'])
        if user_id not in ground_truth:
            ground_truth[user_id] = set()
        ground_truth[user_id].add(movie_id)
    
    print(f"Built ground truth for {len(ground_truth)} users")
    return ground_truth

def prepare_recommendations(
    recommendations: pd.DataFrame,
    k: int = 10
) -> Dict[str, List[str]]:
    """构建推荐列表字典"""
    print(f"Preparing recommendations (top {k})...")
    
    recommendations = recommendations.sort_values(['user_id', 'rank'])
    rec_dict = {}
    for user_id, group in recommendations.groupby('user_id'):
        top_k_movies = group.head(k)['movie_id'].astype(str).tolist()
        rec_dict[user_id] = top_k_movies
    
    print(f"Prepared recommendations for {len(rec_dict)} users")
    return rec_dict

def analyze_results(
    ndcg_scores: Dict[str, float],
    map_scores: Dict[str, float],
    output_dir: str = "metrics_results"
) -> None:

    os.makedirs(output_dir, exist_ok=True)
    
    results_df = pd.DataFrame({
        'user_id': list(ndcg_scores.keys()),
        'ndcg@10': list(ndcg_scores.values()),
        'map@10': list(map_scores.values())
    })
    
    results_df.to_csv(os.path.join(output_dir, "lightgcn_user_metrics.csv"), index=False)
    print(f"User-level metrics saved to {os.path.join(output_dir, 'lightgcn_user_metrics.csv')}")
    
    print("\nMetrics Statistics:")
    print(f"NDCG@10: Mean = {results_df['ndcg@10'].mean():.4f}, Median = {results_df['ndcg@10'].median():.4f}, Std = {results_df['ndcg@10'].std():.4f}")
    print(f"MAP@10: Mean = {results_df['map@10'].mean():.4f}, Median = {results_df['map@10'].median():.4f}, Std = {results_df['map@10'].std():.4f}")
    
    # 可视化部分（全英文标签）
    fig, axes = plt.subplots(1, 2, figsize=(12, 5))
    
    axes[0].hist(results_df['ndcg@10'], bins=20, alpha=0.7)
    axes[0].axvline(results_df['ndcg@10'].mean(), color='r', linestyle='dashed', linewidth=2, label=f"Mean: {results_df['ndcg@10'].mean():.4f}")
    axes[0].set_title('NDCG@10 Distribution')
    axes[0].set_xlabel('Score')
    axes[0].set_ylabel('Number of Users')
    axes[0].legend()
    
    axes[1].hist(results_df['map@10'], bins=20, alpha=0.7)
    axes[1].axvline(results_df['map@10'].mean(), color='r', linestyle='dashed', linewidth=2, label=f"Mean: {results_df['map@10'].mean():.4f}")
    axes[1].set_title('MAP@10 Distribution')  # 修正原代码中的"wget"错误
    axes[1].set_xlabel('Score')
    axes[1].set_ylabel('Number of Users')
    axes[1].legend()
    
    plt.tight_layout()
    plt.savefig(os.path.join(output_dir, "lightgcn_metrics_distribution.png"))
    print(f"Metrics distribution plot saved to {os.path.join(output_dir, 'lightgcn_metrics_distribution.png')}")
    plt.show()

def main():
    # 文件路径配置
    RECOMMENDATIONS_PATH = "recommendation_results/random50_users_lightgcn_detailed.csv"
    RATINGS_PATH = "filtered_data/ratings_filtered.csv"
    USER_PATH = "filtered_data/user_filtered.csv"
    MOVIES_PATH = "filtered_data/movies_filtered.csv"
    K = 10
    
    recommendations, ratings, users, movies = load_data(
        RECOMMENDATIONS_PATH,
        RATINGS_PATH,
        USER_PATH,
        MOVIES_PATH
    )
    
    rec_dict = prepare_recommendations(recommendations, K)
    ground_truth = prepare_ground_truth(ratings, recommendations)
    
    print(f"\nCalculating NDCG@{K} and MAP@{K}...")
    avg_ndcg, avg_map, ndcg_scores, map_scores = evaluate_recommendations(
        rec_dict, 
        ground_truth, 
        K
    )
    
    print(f"\nEvaluation Results:")
    print(f"Average NDCG@{K}: {avg_ndcg:.4f}")
    print(f"Average MAP@{K}: {avg_map:.4f}")
    
    analyze_results(ndcg_scores, map_scores)

if __name__ == "__main__":
    main()